from pyocd.core.helpers import ConnectHelper
import time

class py32f002b:
    PERIPH_BASE = 0x40000000
    AHBPERIPH_BASE = PERIPH_BASE + 0x00020000
    FLASH_R_BASE = AHBPERIPH_BASE + 0x00002000

    def __init__(self,mcu_name = "py32f002bx5"):
        self.mcu_name = mcu_name
        self.session = None
        self.target = None

        self.flashAddr = 0x08000000
        self.flashSize = 0x00006000 # 24KB
        self.uidAddr = 0x1fff0000
        self.optrAddr = 0x1fff0080
        self.sdkrAddr = 0x1fff0084
        self.btcrAddr = 0x1fff0088
        self.wrprAddr = 0x1fff008C

        self.acr = self.FLASH_R_BASE + 0x00
        self.keyr = self.FLASH_R_BASE + 0x08
        self.optkeyr = self.FLASH_R_BASE + 0x0C
        self.sr = self.FLASH_R_BASE + 0x10
        self.cr = self.FLASH_R_BASE + 0x14
        self.sdkr = self.FLASH_R_BASE + 0x24

        self.crDefault = 0xC0000000

    def __enter__(self):
        self.session = ConnectHelper.session_with_chosen_probe(target_override=self.mcu_name)
        self.session.__enter__()
        self.target = self.session.board.target        
        return self
    
    def __exit__(self, exc_type, exc_value, exc_tb):
        if self.session:
            self.session.__exit__(exc_type, exc_value, exc_tb)

    def _readBit(self,addr, pos):
        if not 0 <= pos <= 31:
            raise ValueError("pos must be between 0 and 31")
        value = self.target.read32(addr)
        ret = (value) & (1 << pos) 
        if ret != 0:
            return 1
        return 0
    
    def reset(self):
        self.target.reset_and_halt()
        time.sleep(0.1)

    def read32(self, addr):
        return self.target.read32(addr)

    def write32(self, addr, value):
        self.target.write32(addr, value)

    def writeBit(self, addr, pos, value):
        if value not in [0, 1]:
            raise ValueError("value must be 0 or 1")
        tmp = self.read32(addr)
        if value == 1:
            tmp = tmp | (1 << pos)
        else:
            tmp = tmp & ~(1 << pos)
        self.write32(addr, tmp)
    
    # 写入指定位置的n位
    def writeBits(self, addr, pos_start, pos_end, value):
        if pos_end < pos_start:
            raise ValueError("pos_end must be greater than pos_start")
        if not 0 < pos_end < 31 or not 0 < pos_start < 31:
            raise ValueError("pos_end and pos_start must be between 0 and 31")
        tmp = self.read32(addr)
        mask = (1 << (pos_end - pos_start + 1)) - 1
        tmp = tmp & ~(mask << pos_start)
        tmp = tmp | (value << pos_start)
        print("tmp: 0x%X" % tmp)
        self.write32(addr, tmp)

    def unlockFlash(self):
        LOCK = 31
        if(self._readBit(self.cr, LOCK) == 1):
            self.write32(self.keyr, 0x4567_0123)
            self.write32(self.keyr, 0xCDEF_89AB)
            if(self._readBit(self.cr, LOCK) == 0):
                print("unlockFlash success")
            else:
                print("unlockFlash failed")
        else:
            print("unlockFlash already unlocked")
    
    def lockFlash(self):
        LOCK = 31
        self.writeBit(self.cr, LOCK, 1)
        if(self._readBit(self.cr, LOCK) != 0):
            print("lockFlash success")
        else:
            print("lockFlash failed")

    def unlockOptFlash(self):
        OPTLOCK = 30
        if(self._readBit(self.cr, OPTLOCK) == 1):
            self.write32(self.optkeyr, 0x0819_2A3B)
            self.write32(self.optkeyr, 0x4C5D_6E7F)
            if(self._readBit(self.cr, OPTLOCK) == 0):
                print("unlockOptFlash success")
            else:
                print("unlockOptFlash failed")  
        else:
            print("unlockOptFlash already unlocked")

    def lockOptFlash(self):
        OPTLOCK = 30
        self.writeBit(self.cr, OPTLOCK, 1)
        if(self._readBit(self.cr, OPTLOCK) != 0):
            print("lockOptFlash success")
        else:
            print("lockOptFlash failed")

    def setOption(self, addr, value):
        # 1)用之前描述的步骤，清零 OPTLOCK 位
        self.unlockFlash()
        self.unlockOptFlash()

        # 2)检查 BSY 位，确认没有正在进行的 Flash 操作
        sr = self.read32(self.sr)
        BSY = 16
        while self._readBit(self.sr, BSY) == 1:
            time.sleep(0.1)

        # 3)向选项字节寄存器 FLASH_OPTR/FLASH_SDKR/FLASH_BTCR/FLASH_WRP 写期望的值（1~4 个字）
        self.write32(addr, value)

        # 4)置位 OPTSTRT 位
        OPTSTRT = 17
        self.writeBit(self.cr, OPTSTRT, 1)

        # 5)向 main flash 0x4002 2080 地址写任意 32 位数据（触发正式的写操作），这里写0xffffffff
        self.write32(0x4002_2080, 0xffffffff)

        # 6)等待 BSY 位被清零
        while self._readBit(self.sr, BSY) == 1:
            time.sleep(0.1)

        # 7)等待 EOP 拉高，软件清零
        EOP = 0
        # while self._readBit(self.sr, EOP) != 1:
        #     time.sleep(0.1)
        self.writeBit(self.sr, EOP, 0)
        time.sleep(0.1)
        res = self.read32(addr)
        print("result:addr 0x%X, value 0x%X" % (addr, res))

    def launchOption(self):
        OBLAUNCH = 27
        self.writeBit(self.cr, OBLAUNCH, 1)


    def SetSDKAddress(self, start, end):
        if start > 0b1111 or end > 0b1111:
            print("start or end is out of range")
            return
        self.setOption(self.sdkr, (end << 8) | start)
        self.lockFlash()
        self.lockOptFlash()
        self.launchOption() # 立即生效



    def readUid(self): # 返回4个字节的uid
        uid = []
        for i in range(0,4):
            uid.append(self.target.read32(self.uidAddr + i * 4))
        return uid

    def readOptr(self): # 返回flash_optr
        return self.target.read32(self.optrAddr)

    def readSdkr(self): # 返回flash_sdkr
        return self.target.read32(self.sdkrAddr)

    def readBtcr(self): # 返回flash_btcr
        return self.target.read32(self.btcrAddr)

    def readWrpr(self): # 返回flash_wrpr
        return self.target.read32(self.wrprAddr)

    def printInfo(self):
        uid = self.readUid()
        for i in range(0,4):
            print("uid[%d]: 0x%X" % (i, uid[i]))
        print("optr: 0x%X" % self.readOptr())
        print("sdkr: 0x%X" % self.readSdkr())
        print("btcr: 0x%X" % self.readBtcr())
        print("wrpr: 0x%X" % self.readWrpr())
